import os
import cv2
import numpy as np
from glob import glob
from tqdm import tqdm

CROP_ROOT = "data/crops"
TARGET_SIZE = (256, 256)

def fix_image(path):
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        print(f"无法读取: {path}")
        return
    if img.shape[-2:] == TARGET_SIZE:
        return
    # 灰度/单通道
    if len(img.shape) == 2:
        fixed = np.zeros(TARGET_SIZE, dtype=img.dtype)
        h, w = img.shape
        if h > TARGET_SIZE[0] or w > TARGET_SIZE[1]:
            # 中心裁剪
            start_y = (h - TARGET_SIZE[0]) // 2
            start_x = (w - TARGET_SIZE[1]) // 2
            fixed = img[start_y:start_y+TARGET_SIZE[0], start_x:start_x+TARGET_SIZE[1]]
        else:
            # 居中填充
            start_y = (TARGET_SIZE[0] - h) // 2
            start_x = (TARGET_SIZE[1] - w) // 2
            fixed[start_y:start_y+h, start_x:start_x+w] = img
    else:
        # 多通道
        fixed = np.zeros((TARGET_SIZE[0], TARGET_SIZE[1], img.shape[2]), dtype=img.dtype)
        h, w = img.shape[:2]
        if h > TARGET_SIZE[0] or w > TARGET_SIZE[1]:
            start_y = (h - TARGET_SIZE[0]) // 2
            start_x = (w - TARGET_SIZE[1]) // 2
            fixed = img[start_y:start_y+TARGET_SIZE[0], start_x:start_x+TARGET_SIZE[1], :]
        else:
            start_y = (TARGET_SIZE[0] - h) // 2
            start_x = (TARGET_SIZE[1] - w) // 2
            fixed[start_y:start_y+h, start_x:start_x+w, :] = img
    cv2.imwrite(path, fixed)

def main():
    for root, dirs, files in os.walk(CROP_ROOT):
        for file in files:
            if file.endswith(".tif") or file.endswith(".tiff") or file.endswith(".png"):
                fix_image(os.path.join(root, file))
    print("修复完成！")

if __name__ == "__main__":
    main()